#include "main_control.h"
#include "nvtx3/nvToolsExt.h"

void Main_Control::Step_I() {
  Time_Count TC;
  Time_Count TC_Tot;
  TC.start();
  TC_Tot.start();
  this->print_General();
  // Get Input
  cout << " +++  Step_I +++   " << endl;
  get_info.init(Filename_Input);
  get_info.build();
  string f_dst = "serial_dat/Output/DF.Input";
  File_M.copy(Filename_Input, f_dst);
  // Init coeff
  nvtxRangePush("init_coeff");
  this->init_coeff();
  nvtxRangePop();
  nvtxRangePush("build_basis");
  //exit(0);
  this->build_basis();
  //exit(0);
  nvtxRangePop();
  cout << " -- Finished State & Coeff successfully  -- " << endl;
  //
  int Trans_Test_Flag = -1;
  Trans_Test_Flag = get_info.trans_test.Trans_Test_Flag;
  if (Trans_Test_Flag == 1) {
    TC.start();
    cout << "T-coefficients Start : " << endl;
    Hobra hobra_t(bin_t, cg_x, ninJ_x);
    hobra_t.init_2(bin_t, cg_x, ninJ_x);
    nvtxRangePush("Trans_coeff");
    TransCoeff trans_t(sixJ_x, ninJ_x, hobra_t, jaco_nl, single_t, triple_t);
    ///trans_t.build();
    trans_t.build_OA();
    nvtxRangePop();
    string trans_file = "serial_dat/Output/trans_vec.dat";
    string trans_file_ser = "serial_dat/Output/trans_serial.dat";
    // if (Print_Mid_Flag > 0) {
    // ser_t.save(trans_t, trans_file_ser);
    trans_t.print(trans_file);
    //}
    cout << "T-coefficients Time cost : " << TC.sec() << "s" << endl;
    Check_File CF;
    int flag_diff = -1;
    string trans_file_std = get_info.trans_test.trans_file_std;
    flag_diff = CF.diff_two_file(trans_file, trans_file_std);
    if (flag_diff == 1) {
      cout << " ==  Correct Trans ! == " << endl;
    }
    cout << " Step I total Time cost : " << TC_Tot.min() << " min  or "
         << TC_Tot.hour() << " h on " << MainCom.size() << "nodes" << endl;

    exit(0);
  }
  // --- construct Wave ---
}

void Main_Control::build_basis() {
  Time_Count TC;
  TC.start();
  // Build Single-Paticle states
  cout<<"\t -- Start build basis -- "<<endl;
  if (get_info.sp_info.Type == "Read") {
    single_t.build_T(get_info.sp_info.sp_vec);
  } else if (get_info.sp_info.Type == "E_max") {
    single_t.build_T_E(get_info.sp_info.E_max);
  } else {
    cerr << " Wrong @ build SP " << endl;
    cerr << "\t get_info.sp_info.Type : " << get_info.sp_info.Type << endl;
    exit(0);
  }
  cout<<"\t Test P 1"<<endl;
  single_t.build_Tz();
  single_t.print_Tz();
  single_t.print("serial_dat/Output/SP.dat");
  single_t.print_Tz("serial_dat/Output/SP_Tz.dat");
  ser_t.save(single_t, "serial_dat/Output/SP_serial.dat");
  // Build Triple-Paticle states
  // Get TP info
    cout<<"\t Test P 1"<<endl;
  vector<vector<int>> E_Jac_conf_t;
  vector<int> abc_vec(6);
  int TP_allow_file_num;
  // int TP_E_max;
  vector<vector<vector<int>>> TP_allow_gather;
  vector<int> E_cut_gather;
  vector<int> Same_Particles_gather;
  vector<vector<int>> E_sign_gather;
  TP_allow_file_num = get_info.tp_al_info.size();
  for (int i = 0; i < TP_allow_file_num; i++) {
    vector<vector<int>> TP_allow_vec(2);
    for (int j = 0; j < 2; j++) {
      TP_allow_vec[j].resize(6);
      for (int k = 0; k < 6; k++) {
        TP_allow_vec[j][k] = get_info.tp_al_info[i].tp_al_vec[k][j];
      }
    }
    vector<int> E_sign(6);
    for (int k = 0; k < 6; k++) {
      E_sign[k] = get_info.tp_al_info[i].tp_al_vec[k][2];
    }
    int Same_Particles_Flag = get_info.tp_al_info[i].Same_particle_Flag;
    int E_cut = get_info.tp_al_info[i].E_cut;
    TP_allow_gather.push_back(TP_allow_vec);
    E_cut_gather.push_back(E_cut);
    E_sign_gather.push_back(E_sign);
    Same_Particles_gather.push_back(Same_Particles_Flag);
    if (get_info.tp_info.Type == "abc") {
      single_t.print(TP_allow_vec);
    } else {
      single_t.print();
    }
  }
  // ---- build TP  ----
  cout << "Start build_T" << endl;
  if (get_info.tp_info.Type == "abc") {
    int a_min = get_info.tp_info.abc_vec[0];
    int a_max = get_info.tp_info.abc_vec[1];
    int b_min = get_info.tp_info.abc_vec[2];
    int b_max = get_info.tp_info.abc_vec[3];
    int c_min = get_info.tp_info.abc_vec[4];
    int c_max = get_info.tp_info.abc_vec[5];
    // triple_t.build_T(single_t, a_min, a_max, b_min, b_max, c_min, c_max);
    triple_t.init_Constrain(TP_allow_gather);
    triple_t.build_T_Constrain(single_t, a_min, a_max, b_min, b_max, c_min,
                               c_max);
  } else if (get_info.tp_info.Type == "E_max") {
    int Emax_t = get_info.tp_info.E_max;
    // triple_t.build_T(single_t, Emax_t);
    triple_t.init_Constrain(TP_allow_gather);
    triple_t.build_T_Constrain(single_t, Emax_t);
  } else {
    cerr << " Wrong @ build TP " << endl;
    cerr << "\t get_info.tp_info.Type : " << get_info.tp_info.Type << endl;
    exit(0);
  }
  // cout << "Start build_T test point 1" << endl;

  // triple_t.Select_FI_T(single_t, TP_allow_gather);

  // triple_t.print("serial_dat/Output/TP_before_delete.dat");
  triple_t.Delet_Sort_T();
  // cout << "Finished Delet_Sort_T" << endl;
  triple_t.Delet_conf_T(single_t, E_sign_gather, E_cut_gather,
                        Same_Particles_gather);
  // cout << "Finished Delet_conf_T" << endl;
  triple_t.build_conf_T(single_t, E_sign_gather, E_cut_gather,
                        Same_Particles_gather);
  // cout << "Finished build_conf_T" << endl;
  triple_t.find_E_Jac_conf(E_Jac_conf_t);
  // cout << "Start build_T test point 2" << endl;
  std::ofstream E_conf_file;
  E_conf_file.open("serial_dat/Output/E_conf.Input");
  E_conf_file << E_Jac_conf_t.size() << "\t E_conf Num" << endl;
  for (int i = 0; i < E_Jac_conf_t.size(); i++) {
    E_conf_file << E_Jac_conf_t[i][0] << "\t " << E_Jac_conf_t[i][1]
                << "\t E_F \t E_I" << endl;
  }
  E_conf_file.close();
  // cout << "Start build_T test point 3" << endl;
  triple_t.print_TPconf_info("serial_dat/Output/TPconf.dat",
                             "serial_dat/Output/TPconf_E.dat",
                             "serial_dat/Output/TPconf_E_tot.dat");
  triple_t.sort();
  triple_t.print("serial_dat/Output/TP.dat");
  triple_t.print_channel("serial_dat/Output/TP_channel.dat");
  // ---- build TP_Tz ----
  /*
  cout << "Start build_Tz" << endl;
  if (get_info.tp_info.Type == "abc") {
    int a_min = 2 * get_info.tp_info.abc_vec[0];
    int a_max = 2 * get_info.tp_info.abc_vec[1] + 1;
    int b_min = 2 * get_info.tp_info.abc_vec[2];
    int b_max = 2 * get_info.tp_info.abc_vec[3] + 1;
    int c_min = 2 * get_info.tp_info.abc_vec[4];
    int c_max = 2 * get_info.tp_info.abc_vec[5] + 1;
    triple_t.build_Tz(single_t, a_min, a_max, b_min, b_max, c_min, c_max);
  } else if (get_info.tp_info.Type == "E_max") {
    int Emax_t = get_info.tp_info.E_max;
    triple_t.build_Tz(single_t, Emax_t);
  } else {
    cerr << " Wrong @ build TP_Tz " << endl;
    cerr << "\t get_info.tp_info.Type : " << get_info.tp_info.Type << endl;
    exit(0);
  }
  triple_t.Select_FI_Tz(TP_allow_gather);
  triple_t.Delet_Sort_Tz();
  triple_t.Delet_conf_Tz(single_t, E_sign_gather, E_cut_gather,
                         Same_Particles_gather);
  triple_t.build_conf_Tz(single_t, E_sign_gather, E_cut_gather,
                         Same_Particles_gather);
  triple_t.print_TPconf_info_Tz("serial_dat/Output/TPconf_Tz.dat",
                                "serial_dat/Output/TPconf_E_Tz.dat");
  triple_t.print_Tz("data/TP_Tz.dat");
  triple_t.print_channel_Tz("data/TP_channel_Tz.dat");
  */
  ser_t.save(triple_t, "serial_dat/Output/TP_serial.dat");
  cout << "Finished build_T" << endl;
  //
  // ---- Jacobi ----
  cout << "Start Jacobi " << endl;
  int Emax_JC = triple_t.Emax;
  int J2max_JC = triple_t.J2_max_TP + 2 * triple_t.Emax;
  jaco_b.build_JT(Emax_JC, J2max_JC);
  // cout << "Start Jacobi test point 1" << endl;
  jaco_b.print("serial_dat/Output/Jaco_unselect.dat");
  jaco_b.select_FI_T(triple_t.TPconf_E_tot);
  jaco_b.Delet_Sort_T(triple_t);
  cout << "Start Jacobi test point 2" << endl;
  jaco_b.sort();
  jaco_b.print("serial_dat/Output/Jaco_unselect2.dat");
  jaco_b.print_channel("serial_dat/Output/Jaco_channel_unselect2.dat");
  // -- Antisym --
  cout << " Start Antisym " << endl;
  Hobra hobra_t(bin_t, cg_x, ninJ_x);
  Antisym antisym_t(sixJ_x, ninJ_x, hobra_t, jaco_b);
  cout << " Start Antisym test point 1  " << endl;
  double Anti_test =
      (1 - 2 * antisym_t.mat_cal(jaco_b.state[1][0], jaco_b.state[1][0])) / 3;
  cout << " \t ===+++=== TEST === Anti_test = " << Anti_test << endl;
  antisym_t.cal();
  antisym_t.print("serial_dat/Output/anti_vec.dat");
  jaco_b.print_channel("serial_dat/Output/Jaco_channel.dat");
  jaco_b.select();
  jaco_b.print("serial_dat/Output/Jaco_unselect3.dat");
  cout << "Jacobi State Finished & Antisymmetry Time cost : " << TC.sec() << "s"
       << endl;
  // --  Jacobi_NL  --
  cout << "Start to build Jaco_NL" << endl;
  cout << "\t Triple Emax = " << triple_t.Emax << endl;
  cout << "\t Triple Emax_Final = " << triple_t.Emax_Final;
  cout << "\t Triple Emax_Init = " << triple_t.Emax_Init << endl;
  jaco_nl.build_NL(jaco_b, triple_t);
  cout << "Finished build_NL" << endl;
  jaco_nl.select_FI_T(triple_t);
  cout << "Finished select_FI_T" << endl;
  vector<int> Jaco_used_Flag;
  jaco_nl.Delete_unused_Jaco_channel(jaco_b, Jaco_used_Flag);
  jaco_b.Delete_unused_Jaco_channel(Jaco_used_Flag);
  jaco_nl.sort();
  jaco_b.build_conf_T(triple_t.TPconf_E_tot);
  jaco_b.print_Jacconf("serial_dat/Output/Jacconf.dat");
  jaco_nl.print("serial_dat/Output/Jaco_NL.dat");
  jaco_nl.print_channel("serial_dat/Output/Jaco_NL_channel.dat");
  jaco_b.print("serial_dat/Output/Jaco.dat");
  jaco_b.print_channel("serial_dat/Output/Jaco_channel.dat");
  cout << "Jaco l12max = " << jaco_b.l12max << "\t l3max = " << jaco_b.l3max
       << "\t lmax = " << jaco_b.lmax << endl;
  cout << "Jaco n12max = " << jaco_b.n12max << "\t n3max = " << jaco_b.n3max
       << "\t nmax = " << jaco_b.nmax << endl;
  cout << "Jaco Emax_F = " << jaco_b.Emax_F
       << "\t Jaco Emax_I = " << jaco_b.Emax_I << endl;
  //
  ser_t.save(antisym_t, "serial_dat/Output/anti_serial.dat");
  ser_t.save(jaco_b, "serial_dat/Output/jaco_serial.dat");
  ser_t.save(jaco_nl, "serial_dat/Output/jacoNL_serial.dat");
}

void Main_Control::load_basis() {
  ser_t.load(single_t, "serial_dat/Output/SP_serial.dat");
  if (MainCom.rank() == 0)
    cout << "Finished Load single_t" << endl;
  ser_t.load(triple_t, "serial_dat/Output/TP_serial.dat");
  if (MainCom.rank() == 0)
    cout << "Finished Load triple_t" << endl;
  ser_t.load(jaco_nl, "serial_dat/Output/jacoNL_serial.dat");
  // ser_t.load(antisym_t, "serial_dat/Output/anti_serial.dat");
  ser_t.load(jaco_b, "serial_dat/Output/jaco_serial.dat");
}

void Main_Control::init_coeff() {
  int J_max_cg = get_info.coeff_para.J_max_cg;
  int J_max_sixJ = get_info.coeff_para.J_max_sixJ;
  int n_max = get_info.coeff_para.n_max_bino;
  cout << " init_coeff : J_max_cg = " << J_max_cg
       << "  J_max_sixJ = " << J_max_sixJ << "  n_max = " << n_max << endl;
  // cg_x.init(J_max_cg);
  // sixJ_x.init(J_max_sixJ);
  // ninJ_x.init(5, sixJ_x);
  int n_max_t = 100;
  bin_t.init(n_max_t);

  // ser_t.save(cg_x, "serial_dat/Output/cg_serial.dat");
  // ser_t.save(sixJ_x, "serial_dat/Output/sixJ_serial.dat");
  // ser_t.save(bin_t, "serial_dat/Output/binom_serial.dat");
  // ser_t.save(ninJ_x, "serial_dat/Output/ninJ_serial.dat");
  //hobra_t.init(bin_t, cg_x, ninJ_x);
  cout<<" Finished init coeff"<<endl;
}

void Main_Control::load_coeff() {
  ser_t.load(cg_x, "serial_dat/Output/cg_serial.dat");
  ser_t.load(sixJ_x, "serial_dat/Output/sixJ_serial.dat");
  ser_t.load(bin_t, "serial_dat/Output/binom_serial.dat");
  ser_t.load(ninJ_x, "serial_dat/Output/ninJ_serial.dat");
  // hobra_t.init(bin_t, cg_x, ninJ_x);
}


void Main_Control::print_General() {


    cout<<"  ==============================================================="
            << endl;
    cout<<endl;
    cout<<"       ____             __      ______                     __    "<<endl;
    cout<<"      / __ \\____ ______/ /__   / ________  ________  _____/ /_   "<<endl;
    cout<<"     / / / / __ `/ ___/ //_/  / /_  / __ \\/ ___/ _ \\/ ___/ __/  "<<endl;
    cout<<"    / /_/ / /_/ / /  / .\\    / __/ / /_/ / /  /  __/__  / /_      "<<endl;
    cout<<"   /_____/\\__,_/_/  /_/|_|  /_/    \\____/_/   \\___/____/\\__/  "<<endl;
    //
    cout<<endl;
   cout << "                         version 2.1                         "
        << endl;
   cout << "                     by Yuanzhuo 2020-June                   "
        << endl;
   cout << "    ------------------------------------------------------   "
        << endl;
  Time_Count TC;
  TC.print_time();
  cout << endl;
}


void Main_Control::print_General(std::ofstream &file) {
  file <<"  ==============================================================="
       << endl;
       file<<"       ____             __      ______                     __    "<<endl;
       file<<"      / __ \\____ ______/ /__   / ________  ________  _____/ /_   "<<endl;
       file<<"     / / / / __ `/ ___/ //_/  / /_  / __ \\/ ___/ _ \\/ ___/ __/  "<<endl;
       file<<"    / /_/ / /_/ / /  / .\\    / __/ / /_/ / /  /  __/__  / /_      "<<endl;
       file<<"   /_____/\\__,_/_/  /_/|_|  /_/    \\____/_/   \\___/____/\\__/  "<<endl;
  file << endl;
  file << "    ------------------------------------------------------   "<< endl;
  file << endl;
  Time_Count TC;
  TC.print_time(file);
  file << endl;
}

void Main_Control::print_LECs() {
  cout << "\t ---- LECs : ----" << endl;
  cout << "\t c1 = " << LECs.LEC_C1 << endl;
  cout << "\t c3 = " << LECs.LEC_C3 << endl;
  cout << "\t c4 = " << LECs.LEC_C4 << endl;
  cout << "\t cD = " << LECs.LEC_CD << endl;
  cout << "\t cE = " << LECs.LEC_CE << endl;
  cout << "\t ----------------" << endl;
}
